import os

os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'

import torch
import numpy as np
import matplotlib.pyplot as plt
from digital_freq import digital_freq
from wdm_multiplex import wdm_multiplex
from wdm_demultiplex import wdm_demultiplex
from get_lambda import get_lambda
from read_matlab_struct import read_matlab_struct


def main():
    Signal = {
        'Symbol_number': 2 ** 10,
        'Bits_per_symbol': 4,
        'total_Baud': 42e9,
        'up_sampling_factor': 16,
        'C': 3e8,
        'polarization': 2,
        'lambda_wave': 1550e-12,
        'f_c': 3e8 / 1550e-12,
        'powerdBm_Set': 0,
        'WDM_channel': 10,
        'Ncut': int(np.ceil(100 / 2)),
        'channel_spacing_Hz': 50e9,
    }
    Signal['Sampling_num'] = Signal['Symbol_number'] * Signal['up_sampling_factor']
    Signal['Sampling_rate'] = Signal['total_Baud'] * Signal['up_sampling_factor']
    Signal['powermW_Set'] = 10 ** (Signal['powerdBm_Set'] / 10)
    Signal['channel_spacing_nm'] = Signal['channel_spacing_Hz'] * Signal['lambda_wave'] ** 2 / Signal['C']
    Signal['f'] = digital_freq(Signal)
    Signal['simu_time'] = np.arange(Signal['Sampling_num']) / Signal['Sampling_rate']
    Signal['df_matrix'] = get_lambda(Signal['WDM_channel'], Signal['channel_spacing_Hz'])

    file_path = 'Tx_sym_pulse_shaping.mat'
    struct_name = 'Tx_sym_pulse_shaping'  # Replace with actual struct name in your .mat file

    # Read the data from the .mat file
    X_values, Y_values = read_matlab_struct(file_path, struct_name)

    # 将结构化数组转换为复数数组
    X_complex = X_values['real'] + 1j * X_values['imag']
    Y_complex = Y_values['real'] + 1j * Y_values['imag']

    # 绘制复数数据
    plt.plot(X_complex.real, X_complex.imag, '.', markersize=0.2)
    plt.title('Complex Data Plot')
    plt.xlabel('Real Part')
    plt.ylabel('Imaginary Part')
    plt.grid(True)
    plt.show()

    WDM_sigout_X = wdm_multiplex(Signal, X_values)
    WDM_sigout_Y = wdm_multiplex(Signal, Y_values)
    WDM_sigout = torch.stack([WDM_sigout_X, WDM_sigout_Y], dim=1)
    f = digital_freq(Signal)
    f = np.reshape(f, (Signal['Sampling_num'], 1))
    plt.plot(f, 20 * np.log10(np.abs(np.fft.fftshift(np.fft.fft(WDM_sigout_X.numpy())))), linewidth=1)
    plt.title('WDM')
    plt.xlabel('Frequency (Hz)')
    plt.ylabel('Magnitude (dB)')
    plt.grid(True)
    plt.show()

    WDM_sigout_X, WDM_sigout_Y = wdm_demultiplex(WDM_sigout, Signal)

    # Plotting part
    f = digital_freq(Signal)
    plt.plot(f, 20 * np.log10(np.abs(np.fft.fftshift(np.fft.fft(WDM_sigout_X[:, 0].numpy())))), linewidth=0.5)
    plt.title('WDM')
    plt.xlabel('Frequency (Hz)')
    plt.ylabel('Magnitude (dB)')
    plt.grid(True)
    plt.show()

    # for k in range(Signal['WDM_channel']):
    #     plt.figure()
    #     plt.plot(WDM_sigout_X[:, k].real, WDM_sigout_X[:, k].imag, '.', markersize=0.2)
    #     plt.title(f'Constellation of channel {k + 1} x-pol')
    #     plt.xlabel('Real')
    #     plt.ylabel('Imaginary')
    #     plt.grid(True)
    #     plt.show()


if __name__ == '__main__':
    main()
